Snippets For Python

Data Structures

Segment Tree

class Seg:
    def __init__(self, size, T):
        import math
        self.size = size
        self.h = math.ceil(math.log2(size + 1)) + 1
        self.seg = [0] * (1 << self.h)
        self.seg_zero = 1 << (self.h - 1)
        self.T = T
    
    def update(self, ind, val):
        ind += self.seg_zero
        self.seg[ind] = val
        while ind > 1:
            ind >>= 1
            l = r = ind << 1
            r += 1
            self.seg[ind] = self.T(self.seg[l], self.seg[r])

    def query(self, l, r):
        l += self.seg_zero
        r += self.seg_zero
        ans = 0
        while l <= r:
            cover = l
            dep = 1
            while l % 2 == 0 and cover + (1 << (dep - 1)) <= r:
                l //= 2
                cover += 1 << (dep - 1)
                dep += 1
            ans = self.T(ans, self.seg[l])
            l = cover + 1
        return ans

Disjoint Set, Union find

class DSU:
    def __init__(self, N):
        self.rank = [1] * (N + 1)
        self.dsu = [_ for _ in range(N + 1)]
    
    def find(self, child):
        if self.dsu[child] == child:
            return child
        else:
            r = self.find(self.dsu[child])
            self.dsu[child] = r
            return self.dsu[child]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return
        if self.rank[x] > self.rank[y]:
            self.dsu[y] = x
            self.rank[x] += self.rank[y]
        else:
            self.dsu[x] = y
            self.rank[y] += self.rank[x]

Math

Matrix Operation

class MatrixOP:
    def mul(self, a, b, MOD):
        N = len(a)
        return [[sum(a[i][k] * b[k][j] % MOD for k in range(N)) % MOD for j in range(N)] for i in range(N)]

    def diag(self, N):
        B = [[0] * N for _ in range(N)]
        for i in range(N):
            B[i][i] = 1
        return B

    def matpow(self, A, T, MOD):
        B = self.diag(len(A))
        while T:
            if T % 2:
                B = self.mul(A, B, MOD)
                T -= 1
            A = self.mul(A, A, MOD)
            T //= 2
        return B